#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@163.com
@desc: 
"""
import argparse
from pathlib import Path
from typing import Iterable

import cv2
import numpy as np
import tensorflow as tf
from dltools.train.tools import run_detection
from dltools.utils.log import get_console_logger
from object_detection.utils import visualization_utils as vis_utils

from utils.evaluator import Evaluator
from utils.file_read import GoodsReader


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--image-root', type=str, default='', dest='root',
        help='The directory where the image data and the annotation is stored.'
             '保存图像与标记数据的根目录, 图像与其相应标记应当在同一级目录。')
    parser.add_argument(
        '--output-root', type=str, default='', dest='output',
        help='the root for all the output.'
             '输出结果的根目录')
    parser.add_argument('--model-dir', type=str, default='', dest='model_dir',
                        help='指定模型的路径')
    parser.add_argument('--cpu-only', type=bool, default=False, dest='cpu_only',
                        help='指定是否使用 CPU')
    parser.add_argument('--score', type=float, default=0.5, dest='score',
                        help='显示的置信度阈值')
    return parser.parse_args()


class GoodsEvaluator(Evaluator):

    def __init__(self, param, logger=None):
        super(GoodsEvaluator, self).__init__(param, logger=logger)
        self.category_index = {}
        for idx in range(30):
            self.category_index[idx + 1] = {'id': idx + 1,
                                            'name': 'package%d' % idx}

    def _detect(self, sess, image, xml, output):
        img_arr = cv2.imread(image)
        cv2.imwrite(output + '_origin.jpg', img_arr)
        img_arr = cv2.cvtColor(img_arr, cv2.COLOR_BGR2RGB)
        img = np.expand_dims(img_arr, axis=0).astype(np.float32)
        boxes, classes, scores = run_detection(sess, self.models, img)
        vis_utils.visualize_boxes_and_labels_on_image_array(
            img_arr, boxes, classes, scores, self.category_index,
            use_normalized_coordinates=True,
            min_score_thresh=self.params['score'],
            line_thickness=5)
        img_arr = cv2.cvtColor(img_arr, cv2.COLOR_RGB2BGR)
        cv2.imwrite(output + '_detected.jpg', img_arr)

    def _update(self, dataset, output):
        assert isinstance(dataset, Iterable)
        output = Path(output)
        if not output.exists():
            output.mkdir(parents=True)
        with self.models.as_default():
            with tf.Session(graph=self.models, config=self.config) as sess:
                with tf.device(self.device):
                    for idx, data in enumerate(dataset):
                        if idx % self.params['display'] == 0:
                            self._logger.info(
                                'Detecting {}-th Data object ...'.format(idx))
                        self._detect(sess, **data,
                                     output=str(output / str(idx)))


if __name__ == '__main__':
    ARGS = parse_args()
    _logger = get_console_logger('GoodsEvaluator')
    _data_reader = GoodsReader(ARGS.root, _logger)
    PARAMS = {'model_dir': ARGS.model_dir, 'root': ARGS.root,
              'score': ARGS.score, 'display': 1, 'cpu_only': ARGS.cpu_only}
    detector = GoodsEvaluator(PARAMS, _logger)
    detector.update(_data_reader, ARGS.output)
